from socket import socket, AF_INET, SO_REUSEADDR, SOCK_STREAM, SOL_SOCKET, timeout
from NNserver.myThread import FC_Para
from Data_Pack import myPack, myRecv
from NNserver.myDset import CFP_Random, CFP_Random, CFP_Valid
from torch.utils.data import DataLoader
import torch
import numpy as np
import random
import os


RCFP_BSIZE = 50
VCFP_BSIZE = 35


class myServer():
    def __init__(self, args):
        # 随机数
        if args.seed != 0:
            torch.manual_seed(args.seed)
            torch.cuda.manual_seed_all(args.seed)
            np.random.seed(args.seed)
            random.seed(args.seed)
            torch.use_deterministic_algorithms(True)
            torch.backends.cudnn.benchmark = False
            os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
            os.environ['PYTHONHASHSEED'] = str(args.seed)
        # 构造工作线程
        self.myPa = FC_Para()
        self.work_on = False
        # 构造连接套接字
        self.sc_s = socket(AF_INET, SOCK_STREAM)
        self.wsec_s = 120.0
        self.sc_c = None
        self.wsec_c = 120.0
        self.addr_c = None
        # 构造指令-函数映射关系
        self.cmd_dic = {
            'close': self.Work_close,
            'loaddset': self.Work_ldset,
            'samget': self.Work_samget,
            'slwfile': self.Work_slwfile,
            'nnt': self.Work_nnt,
            'nns': self.Work_nns,
            'nnvcfp': self.Work_nnvcfp,
            'nngs': self.Work_nngs,
            'nngd': self.Work_nngd
        }
        self.cmd_fun = None
        self.tcp_text = None
        # 绑定地址，连接客户端
        self.Con_wait(args)

    def Con_wait(self, args):
        self.sc_s.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
        self.sc_s.settimeout(self.wsec_s)
        try:
            self.sc_s.bind((args.host, args.port))
        except Exception as e:
            print(e)
            return
        self.sc_s.listen(1)
        print('--正在监听：{}'.format(self.sc_s.getsockname()))
        try:
            self.sc_c, self.addr_c = self.sc_s.accept()
        except Exception as e:
            print(e)
        else:
            print('收到连接：{}'.format(self.addr_c))
            self.sc_c.settimeout(self.wsec_c)
            self.work_on = True

    def Work(self):
        print('--正在等待客户端指令')
        try:
            self.tcp_text = myRecv(self.sc_c)
            self.cmd_fun = self.cmd_dic[self.tcp_text['cmd']]
        except timeout as e:
            print(e)
            return
        except Exception as e:
            print(type(e), e)
            self.sc_c.close()
            self.work_on = False
            return
        print('收到指令：{}'.format(self.tcp_text['cmd']))
        self.cmd_fun()

    def Work_close(self):
        self.sc_c.close()
        self.work_on = False

    def seed_worker(worker_id):
        worker_seed = torch.initial_seed() % 2**32
        np.random.seed(worker_seed)
        random.seed(worker_seed)

    def Work_ldset(self):
        self.myPa.cfp_path = self.tcp_text['cfp_path']
        print('数据集加载中')
        try:
            self.myPa.dset_rcfp = CFP_Random(self.myPa.cfp_path)
            self.myPa.dset_vcfp = CFP_Valid(self.myPa.cfp_path)
        except Exception as e:
            print(e)
            self.myPa.dset_on = False
            self.tcp_text = {'dset_on': False}
            print('数据集加载失败！')
        else:
            self.myPa.dload_rcfp = DataLoader(self.myPa.dset_rcfp, batch_size=RCFP_BSIZE, shuffle=False,
                                                    drop_last=False, worker_init_fn=self.seed_worker)
            self.myPa.dload_vcfp = DataLoader(self.myPa.dset_vcfp, batch_size=VCFP_BSIZE, shuffle=False,
                                                    drop_last=False, worker_init_fn=self.seed_worker)
            self.myPa.dset_on = True
            self.tcp_text = {'dset_on': True,
                            'rcfp_len': len(self.myPa.dset_rcfp)}
            print('数据集加载成功！')
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_samget(self):
        self.myPa.sam_idx = self.tcp_text['sam_idx']
        print('正在获取样本图')
        self.myPa.sam_tsr = self.myPa.dset_rcfp[self.myPa.sam_idx]
        self.myPa.sam_tsr = torch.cat((self.myPa.sam_tsr[0],
                                    self.myPa.sam_tsr[1],
                                    self.myPa.sam_tsr[2]), dim=2)
        self.myPa.sam_img = self.myPa.sam_trans(self.myPa.sam_tsr)
        self.tcp_text = {'sam_img': self.myPa.sam_img,
                        'sam_a_idx': self.myPa.dset_rcfp.a_idx,
                        'sam_p_idx': self.myPa.dset_rcfp.p_idx,
                        'sam_n_idx': self.myPa.dset_rcfp.n_idx,
                        'sam_p_str': self.myPa.dset_rcfp.p_str,
                        'sam_n_str': self.myPa.dset_rcfp.n_str}
        print('样本图获取完成！')
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_slwfile(self):
        self.myPa.nn_wsave = self.tcp_text['nn_wsave']
        if self.myPa.nn_wsave:
            print('正在获取权重文件')
            self.myPa.nn_wfile = {'s_net': self.myPa.nn_net.state_dict(),
                                    's_opt': self.myPa.nn_opt.state_dict(),
                                    'f_thres': self.myPa.nn_net.module.f_thres}
            self.tcp_text['nn_wfile'] = self.myPa.nn_wfile
            print('权重文件获取完成！')
            self.sc_c.sendall(myPack(self.tcp_text))
        else:
            print('正在载入权重文件')
            self.myPa.nn_wfile = self.tcp_text['nn_wfile']
            self.tcp_text = {'nn_wfok': True}
            self.sc_c.sendall(myPack(self.tcp_text))
            self.myPa.nn_net.load_state_dict(self.myPa.nn_wfile['s_net'])
            self.myPa.nn_opt.load_state_dict(self.myPa.nn_wfile['s_opt'])
            self.myPa.nn_net.module.f_thres = self.myPa.nn_wfile['f_thres']
            print('权重文件载入完成！')
        

    def Work_nnt(self):
        nnt_ok = not self.myPa.nn_workon
        if nnt_ok:
            print('神经网络开始训练')
            self.myPa.nnt_epoch = self.tcp_text['nnt_epoch']
            self.myPa.nn_workon = True
            self.myPa.nn_mode = 1
            self.myPa.Th_run(1)
        else:
            print('神经网络正在运作')
        self.tcp_text = {'nnt_ok': nnt_ok,
                        'nn_workon': self.myPa.nn_workon,
                        'nn_mode': self.myPa.nn_mode}
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nns(self):
        self.tcp_text = {'nn_mode': self.myPa.nn_mode}
        if self.myPa.nn_workon:
            print('神经网络正在中止')
            self.tcp_text['nns_ok'] = True
            self.tcp_text['nn_workon'] = False
            self.myPa.nn_sn = True
        else:
            self.tcp_text['nns_ok'] = False
            self.tcp_text['nn_workon'] = self.myPa.nn_workon
            print('神经网络未在运行')
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nnvcfp(self):
        nnvcfp_ok = not self.myPa.nn_workon
        if nnvcfp_ok:
            print('神经网络开始CFP验证')
            self.myPa.nn_workon = True
            self.myPa.nn_mode = 2
            self.myPa.Th_run(2)
        else:
            print('神经网络正在运作')
        self.tcp_text = {'nnvcfp_ok': nnvcfp_ok,
                        'nn_workon': self.myPa.nn_workon,
                        'nn_mode': self.myPa.nn_mode}
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nngs(self):
        self.tcp_text = {'nn_workon': self.myPa.nn_workon,
                        'nn_mode': self.myPa.nn_mode,
                        'nn_nowepoch': self.myPa.nn_nowepoch,
                        'nn_nowbatch': self.myPa.nn_nowbatch}
        self.sc_c.sendall(myPack(self.tcp_text))

    def Work_nngd(self):
        self.tcp_text = {'nn_aok': self.myPa.nn_aok,
                        'nn_lok': self.myPa.nn_lok,
                        'a_depoch': self.myPa.nnt_epoch}
        if self.myPa.nn_aok:
            self.tcp_text['a_num'] = self.myPa.a_num
            self.tcp_text['a_cor'] = self.myPa.a_cor
            self.tcp_text['a_rate'] = self.myPa.a_rate
        if self.myPa.nn_lok:
            self.tcp_text['a_losses'] = self.myPa.a_losses
        self.sc_c.sendall(myPack(self.tcp_text))
